#PBAR_TestZspecBasicHighZFinder_Single.py
#
#
#

import glob, os
import numpy as np
import PBAR_Zspec, PBAR_Cargo
import matplotlib.pyplot as plt
import copy

# image processing stuff
from skimage.filter import canny
from scipy import ndimage
from skimage import morphology

from skimage.filter import threshold_otsu
from skimage.segmentation import clear_border
from skimage.morphology import label, closing, square
from skimage.measure import regionprops

#from skimage.color import label2rgb
import PIL
from PIL import Image
reload(PBAR_Zspec)


(goodZspecMask, badZspecMask, goodZspecNameList, badZspecNameList, \
        goodZspecIndices, badZspecIndices) = PBAR_Zspec.ZspecDetectorLists()

# energy centers
energy = np.array([7.50E+00,1.05E+01,1.35E+01,1.65E+01,1.95E+01,2.25E+01,2.55E+01, \
    2.85E+01,3.15E+01,3.45E+01,3.75E+01,4.05E+01,4.35E+01,4.65E+01,4.95E+01,\
    5.25E+01,5.55E+01,5.85E+01,6.15E+01,6.45E+01,6.75E+01,7.05E+01,7.35E+01, \
    7.65E+01,7.95E+01,8.25E+01,8.55E+01,8.85E+01,9.15E+01,9.45E+01,9.75E+01,1.01E+02])
    
# get linear discriminant coefficients, polynomial fits
(wBest, pfitmean, pfitsigma) = PBAR_Zspec.GetZspecDiscrimValues()

# explore making mask
dataDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansStandardWidth'
cargoConfigBaseDirList = glob.glob(os.path.join(dataDir, '*'))

# round-about way of selecting the desired dataset 
for (dirIndex, cargoConfigBaseDir) in enumerate(cargoConfigBaseDirList):
    subDirList = glob.glob(os.path.join(cargoConfigBaseDir, '*'))
    for (subDirIndex, dataPath) in enumerate(subDirList):        
        print(dataPath)
        a, b = os.path.split(dataPath)
        a, c = os.path.split(a)
        if b == 'A':
            break
    if c == '9':
        break

temp = glob.glob(os.path.join(dataPath, '*cargoimage.npy'))
fullFilenameCargo = temp[0]
(a, filenameCargo) = os.path.split(fullFilenameCargo)
# read in the image
datCargoStandardWidth = np.load(fullFilenameCargo)

# find the zspec scan number
temp = glob.glob(os.path.join(dataPath, '*.npy'))
# if file doesn't exist skip this set completely
a, filename = os.path.split(temp[0])
zspecScanNumber = filename[0:4]
filenameZspec = filename
fullFilenameZspec = os.path.join(dataPath, filename)

# read in the markerfiles
fullFilenameMarker = os.path.join(dataPath, fullFilenameCargo.replace('cargoimage.npy', 'cargomarker'))
markerList = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)

# calculate discrimination stuff
datZspecStandardWidth = np.load(fullFilenameZspec)
discrim = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecStandardWidth, energy)   

# Calculate the dist, linear combination
# make features array (num features X time slices X number of detectors)
featuresMatrix = np.zeros((3, discrim['binMean'].shape[0], discrim['binMean'].shape[1]))
featuresMatrix[0,:,:] = discrim['binMean']
featuresMatrix[1,:,:] = discrim['binSTD']
featuresMatrix[2,:,:] = discrim['multibin_20_ratio']

discrim['dist0']  = featuresMatrix[0,:,:] * wBest[0] + featuresMatrix[1,:,:] * wBest[1] + featuresMatrix[2,:,:] * wBest[2]
discrim['discrim'] =  discrim['dist0'] - np.polyval(pfitmean, discrim['count'])

# Reduce the size of the zspec image by summing adjacent bins
newWidth = 600
newHeight = 136
(datZspecSmall, discrimSmall) = PBAR_Zspec.ZspecBasicReduceSize(datZspecStandardWidth, energy, newWidth)

# Modify the marker files
multiplier = newWidth/float(datZspecStandardWidth.shape[0])
offset = 0.0
markerSmallList = PBAR_Zspec.ModifyMarkersXPosition(markerList, multiplier, offset)

##################
##  MAKE MASKs  ##    
xCargoInternalBounds = np.array((38, 550))
yCargoInternalBounds = np.array((13, 106))
cargoCountRange = np.array([0, 0.5e8]) # only consider parts of the zspec image that have rad below this value

# Make reduced size cargo image
datCargoSmall = PBAR_Zspec.CargoReduceSize(datCargoStandardWidth, newHeight, newWidth)
datCargoSmallMask = PBAR_Zspec.CreateCargoMask(datCargoSmall, cargoCountRange, xCargoInternalBounds, yCargoInternalBounds)


## Look for high z stuff

# make different size windows
windowList = []
windowList.append(np.array((2, 2)))
windowList.append(np.array((4, 4)))
windowList.append(np.array((6, 6)))

discrimThreshold = 5.0
fractionPixels = 1.0
mask = copy.copy(datCargoSmallMask)

# look for potential stuff
potential = PBAR_Zspec.BasicScanZspecHighZFinder(discrimSmall, datZspecSmall, mask, energy, windowList, discrimThreshold, fractionPixels)


###########
## PLOTS ##
###########

# REDUCED SIZED IMAGES



# cargo image, Small with mask


plotMarkers = True
plotPotential = True
plt.figure()
plt.grid()

mask = ( discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & datCargoSmallMask

temp  = copy.copy(discrimSmall['discrim'])
cutt = ~np.isnan(temp) & ~np.isinf(temp)
temp[~cutt] = min(temp[cutt])

temp = discrimSmall['discrim'] * mask
minValue = temp[mask].min()
temp[~mask] = minValue

plt.imshow(datCargoSmall.T, interpolation = 'nearest', aspect='auto', cmap = plt.cm.Greys_r)
#plt.imshow(discrimSmall['discrim'].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)

plt.colorbar()
plt.ylabel('Detector Number')
plt.xlabel('Time Slice')
plt.title(fullFilenameZspec + ', Small')

plt.axis((0, 600, 136, 0))

if plotMarkers:
    marker = markerSmallList
    for i, mark in enumerate( marker):
        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
        plt.plot(x, y, 'r')
        plt.plot(mark['x'], mark['y'], 'xr')
        if 'left' in mark:
            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
        if 'right' in mark:
            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)

if plotPotential:
    for (potIndex, pot) in enumerate(potential):
#        plt.plot(pot['boxCenter'][0], pot['boxCenter'][1], 'xb', markersize = 14)
        plt.scatter(pot['boxCenter'][0], pot['boxCenter'][1], color = 'b', \
            s = 50 * pot['window'].mean(), linewidth = pot['window'].mean()/2.0, edgecolor = 'blue', facecolor = 'blue', marker = 'x', alpha = 0.5)



# Zspec image, Small with mask, discrim

climit = (discrimThreshold, 50)

plotMarkers = True
plotPotential = True
plt.figure()
plt.grid()

mask = ( discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & datCargoSmallMask

temp  = copy.copy(discrimSmall['discrim'])
cutt = ~np.isnan(temp) & ~np.isinf(temp)
temp[~cutt] = min(temp[cutt])

temp = discrimSmall['discrim'] * mask
minValue = temp[mask].min()
temp[~mask] = minValue

plt.imshow(temp.T, interpolation = 'nearest', aspect='auto', cmap = plt.cm.Greys_r)
#plt.imshow(discrimSmall['discrim'].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)

plt.colorbar()
plt.ylabel('Detector Number')
plt.xlabel('Time Slice')
plt.title(fullFilenameZspec + ', Small')

plt.axis((0, 600, 136, 0))
plt.clim(minValue, 50)
if plotMarkers:
    marker = markerSmallList
    for i, mark in enumerate( marker):
        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
        plt.plot(x, y, 'r')
        plt.plot(mark['x'], mark['y'], 'xr')
        if 'left' in mark:
            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
        if 'right' in mark:
            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)

if plotPotential:
    for (potIndex, pot) in enumerate(potential):
#        plt.plot(pot['boxCenter'][0], pot['boxCenter'][1], 'xb', markersize = 14)
        plt.scatter(pot['boxCenter'][0], pot['boxCenter'][1], color = 'b', \
            s = 50 * pot['window'].mean(), linewidth = pot['window'].mean()/2.0, edgecolor = 'blue', facecolor = 'blue', marker = 'x', alpha = 0.5)




# Zspec small, scatter plot

removeOffset = True

plt.figure()
plt.grid()

feature1 = 'count'
feature2 = 'discrim'

mask = (discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & datCargoSmallMask

mask = mask[:,goodZspecIndices]

x1 = discrimSmall[feature1][:,goodZspecIndices][mask].flatten()
y1 = discrimSmall[feature2][:,goodZspecIndices][mask].flatten()

x1 = discrimSmall[feature1][:,goodZspecIndices].flatten()
y1 = discrimSmall[feature2][:,goodZspecIndices].flatten()

plt.plot(x1, y1, '.k', alpha = 0.15, markersize = 10, label = 'In Mask')

marker = markerSmallList
for i, mark in enumerate(marker):
    x_range = np.array((mark['rec_left'], mark['rec_right'])) # left < right
    y_range = np.array((mark['rec_top'], mark['rec_bottom'])) # top  < bottom
    xarray = np.arange(discrimSmall[feature1].shape[0])
    yarray = np.arange(discrimSmall[feature1].shape[1])
    xcut = (xarray > x_range[0]) & (xarray < x_range[1])
    ycut = (yarray > y_range[0]) & (yarray < y_range[1])
    x = discrimSmall[feature1][:,goodZspecMask & ycut][xcut,:].flatten()
    y = discrimSmall[feature2][:,goodZspecMask & ycut][xcut,:].flatten()

    try:
        # high density stuff
        if mark['target'][0] == 'S':
            plt.plot(x, y, 'db', markersize = 8, alpha  = 0.3, label = mark['target'])
        elif mark['target'][0] == 'W':
            plt.plot(x, y, 'vb', markersize = 8, alpha  = 0.3, label = mark['target'])
        elif mark['target'][0] == 'D':
            plt.plot(x, y, 'sb', markersize = 8, alpha  = 0.3, label = mark['target'])
        elif mark['target'][0] == 'P':
            plt.plot(x, y, '*b', markersize = 8, alpha  = 0.3, label = mark['target'])
        elif mark['target'][0] == 'F':# low density stuff
            plt.plot(x, y, 'or', markersize = 8, alpha  = 0.3, label = mark['target'])
        else:
            plt.plot(x, y, 'om', markersize = 8, alpha  = 0.3, label = mark['target'])
    except:
        plt.plot(x, y, 'om', markersize = 8, alpha  = 0.3, label = mark['target'])

if plotPotential:
    for (potIndex, pot) in enumerate(potential):
        plt.scatter(pot['discrim']['count'], pot['discrim']['discrim'], color = 'b', \
        s = 100 * pot['window'].mean(), linewidth = pot['window'].mean(), edgecolor = 'blue', facecolor = 'blue', marker = 'x', alpha = 0.5)

plt.legend()



#
## Zspec image, Small
#index = 0
#plotMarkers = True
#    
#plt.figure()
#plt.grid()
#
##mask = ( discrimSmall['discrim'] > 0 ) & (discrimSmall['discrim'] < 50) & masksSmallList[index]
#
#mask = ( discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & masksSmallList[index]
#
#temp = discrimSmall['discrim'] * mask
#minValue = temp[mask].min()
#temp[~mask] = minValue
#
#plt.imshow(temp.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#plt.colorbar()
#plt.ylabel('Detector Number')
#plt.xlabel('Time Slice')
#plt.title(fullFilenameZspec + ', Small')
#
#plt.axis((0, 600, 136, 0))
#
#if plotMarkers:
#    marker = markerSmallList
#    for i, mark in enumerate( marker):
#        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
#        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
#        plt.plot(x, y, 'r')
#        plt.plot(mark['x'], mark['y'], 'xr')
#        if 'left' in mark:
#            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#        if 'right' in mark:
#            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
#        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
#
#
### PLOT FILTERED MASK
#
#index = 3
#plotMarkers = True
#    
#plt.figure()
#plt.grid()
#
#plt.imshow(masksSmallList[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#plt.colorbar()
#plt.ylabel('Detector Number')
#plt.xlabel('Time Slice')
#
#plt.axis((0, 600, 136, 0))
#
#if plotMarkers:
#    marker = markerSmallList
#    for (i, mark) in enumerate(marker):
#        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
#        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
#        plt.plot(x, y, 'r')
#        plt.plot(mark['x'], mark['y'], 'xr')
#        if 'left' in mark:
#            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#        if 'right' in mark:
#            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
#        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
#
#plt.title('%s, %s' %(filenameCargo, masksSmallNameList[index]))
#
#
## Cargo image, standard size
#plotMarkers = True
#    
#plt.figure()
#plt.grid()
#
#plt.imshow(datCargoStandardWidth.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#plt.colorbar()
#plt.ylabel('Detector Number')
#plt.xlabel('Time Slice')
#plt.title(fullFilenameCargo)
#
#plt.axis((0, 2000, 136, 0))
#
#if plotMarkers:
#    marker = markerList
#    for i, mark in enumerate( marker):
#        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
#        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
#        plt.plot(x, y, 'r')
#        plt.plot(mark['x'], mark['y'], 'xr')
#        if 'left' in mark:
#            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#        if 'right' in mark:
#            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
#        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
#
#
#
## Zspec image, standard size
#
#plotMarkers = True
#    
#plt.figure()
#plt.grid()
#
#mask = ( discrim['discrim'] > 10 ) & (discrim['discrim'] < 50)
#
#plt.imshow((discrim['discrim'] * mask).T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#plt.colorbar()
#plt.ylabel('Detector Number')
#plt.xlabel('Time Slice')
#plt.title(fullFilenameZspec)
#
#plt.axis((0, 2000, 136, 0))
#
#if plotMarkers:
#    marker = markerList
#    for i, mark in enumerate( marker):
#        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
#        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
#        plt.plot(x, y, 'r')
#        plt.plot(mark['x'], mark['y'], 'xr')
#        if 'left' in mark:
#            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#        if 'right' in mark:
#            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
#        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
#
#
## image mask
#plotMarkers = True
#    
#plt.figure()
#plt.grid()
#
#plt.imshow(cargoZspecMaskStandardWidth.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#plt.colorbar()
#plt.ylabel('Detector Number')
#plt.xlabel('Time Slice')
#plt.title(fullFilenameCargo)
#
#plt.axis((0, 2000, 136, 0))
#
#if plotMarkers:
#    marker = markerList
#    for i, mark in enumerate( marker):
#        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
#        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
#        plt.plot(x, y, 'r')
#        plt.plot(mark['x'], mark['y'], 'xr')
#        if 'left' in mark:
#            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#        if 'right' in mark:
#            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
#        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
#
#
## image with masked points only
#
#plotMarkers = True
#    
#plt.figure()
#plt.grid()
#
#plt.imshow(cargoZspecMaskStandardWidth.T * datCargoStandardWidth.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#plt.colorbar()
#plt.ylabel('Detector Number')
#plt.xlabel('Time Slice')
#plt.title(fullFilenameCargo)
#
#plt.axis((0, 2000, 136, 0))
#
#if plotMarkers:
#    marker = markerList
#    for i, mark in enumerate( marker):
#        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
#        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
#        plt.plot(x, y, 'r')
#        plt.plot(mark['x'], mark['y'], 'xr')
#        if 'left' in mark:
#            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#        if 'right' in mark:
#            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
#        plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)


#
#masksSmallList = []
#masksSmallNameList = []
#
##1) smaller mask
##remove cargo edges
##remove small objects
##closing
#
#datCargoSmallMaskCleaned = copy.copy(datCargoSmallMask)
#
#xCargoInternalBounds = np.array((38, 550))
#yCargoInternalBounds = np.array((13, 106))
#
#datCargoSmallMaskCleaned[0:xCargoInternalBounds[0],:] = False
#datCargoSmallMaskCleaned[xCargoInternalBounds[1]:,:] = False
#datCargoSmallMaskCleaned[:,0:yCargoInternalBounds[0]] = False
#datCargoSmallMaskCleaned[:,yCargoInternalBounds[1]:]  = False
#
#datCargoSmallMaskCleaned = morphology.remove_small_objects(datCargoSmallMaskCleaned, 3)
#datCargoSmallMaskCleaned = closing(datCargoSmallMaskCleaned, square(2))
#
#masksSmallList.append(datCargoSmallMaskCleaned.astype(bool))
#masksSmallNameList.append('Small, remove small objects 3, closed square(2)')
#
#
##2) smaller mask
##remove cargo edges
##remove small objects
#
#datCargoSmallMaskCleaned = copy.copy(datCargoSmallMask)
#
#xCargoInternalBounds = np.array((38, 550))
#yCargoInternalBounds = np.array((13, 106))
#
#datCargoSmallMaskCleaned[0:xCargoInternalBounds[0],:] = False
#datCargoSmallMaskCleaned[xCargoInternalBounds[1]:,:] = False
#datCargoSmallMaskCleaned[:,0:yCargoInternalBounds[0]] = False
#datCargoSmallMaskCleaned[:,yCargoInternalBounds[1]:]  = False
#
#datCargoSmallMaskCleaned = morphology.remove_small_objects(datCargoSmallMaskCleaned, 3)
#
#masksSmallList.append(datCargoSmallMaskCleaned.astype(bool))
#masksSmallNameList.append('Small, remove small objects 3')
#
#
##3) smaller mask
##remove cargo edges
##remove small objects 2
##closing
#
#datCargoSmallMaskCleaned = copy.copy(datCargoSmallMask)
#
#xCargoInternalBounds = np.array((38, 550))
#yCargoInternalBounds = np.array((13, 106))
#
#datCargoSmallMaskCleaned[0:xCargoInternalBounds[0],:] = False
#datCargoSmallMaskCleaned[xCargoInternalBounds[1]:,:] = False
#datCargoSmallMaskCleaned[:,0:yCargoInternalBounds[0]] = False
#datCargoSmallMaskCleaned[:,yCargoInternalBounds[1]:]  = False
#
#datCargoSmallMaskCleaned = morphology.remove_small_objects(datCargoSmallMaskCleaned, 2)
#
#masksSmallList.append(datCargoSmallMaskCleaned.astype(bool))
#masksSmallNameList.append('Small, remove small objects 2')
#
##4) mask using the zspec
#
#datZspecSmallMask = discrimSmall['discrim'] > 0
#
#datZspecSmallMaskCleaned = copy.copy(datZspecSmallMask)
#
#xCargoInternalBounds = np.array((38, 550))
#yCargoInternalBounds = np.array((13, 106))
#
#datZspecSmallMaskCleaned[0:xCargoInternalBounds[0],:] = False
#datZspecSmallMaskCleaned[xCargoInternalBounds[1]:,:] = False
#datZspecSmallMaskCleaned[:,0:yCargoInternalBounds[0]] = False
#datZspecSmallMaskCleaned[:,yCargoInternalBounds[1]:]  = False
#
#masksSmallList.append(datZspecSmallMaskCleaned.astype(bool))
#masksSmallNameList.append('Zspec Small, threshold 0')
#
#
##5) mask using the zspec
#
#datZspecSmallMask = discrimSmall['discrim'] > 0
#
#datZspecSmallMaskCleaned = copy.copy(datZspecSmallMask)
#
#xCargoInternalBounds = np.array((38, 550))
#yCargoInternalBounds = np.array((13, 106))
#
#datZspecSmallMaskCleaned[0:xCargoInternalBounds[0],:] = False
#datZspecSmallMaskCleaned[xCargoInternalBounds[1]:,:] = False
#datZspecSmallMaskCleaned[:,0:yCargoInternalBounds[0]] = False
#datZspecSmallMaskCleaned[:,yCargoInternalBounds[1]:]  = False
#
#datZspecSmallMaskCleaned = morphology.remove_small_objects(datZspecSmallMaskCleaned, 2)
#
#masksSmallList.append(datZspecSmallMaskCleaned.astype(bool))
#masksSmallNameList.append('Zspec Small, threshold 0, remove small objects 2')

